#
#
# TODO: write the code that actually does something
# TODO: build some ast
#       TODO: do some variable renaming when inlining
# TODO: write some ast matching code, and some "logic variables" code for matching
# TODO: after building the ast of two composed (real) functions,
#       store it in the __ast__ variable so that it can be further composed
# TODO: write some code to express tree patterns by example,
#       using "compiler.parse()" and then extracting the resulting tree
#       (need to beware of ... something)
# TODO: write some constant folding, and assignment shortcut-ing

__all__ = ['compose', 'optimize']

import compiler,copy
from compiler.ast import Return,Assign,Name,AssName,CallFunc,Function

def compose(downstream,upstream):
    composition = ComposedTransducer(downstream,upstream)
    compositions.append(composition)
    composed_functions.add(downstream)
    composed_functions.add(upstream)
    return composition

def optimize():
    modules = {}
    for func in composed_functions:
        if type(func) == ComposedTransducer:
            continue # this avoids doing something dumb like parsing ourself
        
        try:
            modules[func.__module__].add(func)
        except KeyError:
            modules[func.__module__] = set([func])
    
    import sys    
    for name,functions in modules.iteritems():
        module = sys.modules[name]
        filename = module.__file__
        if filename.endswith('.pyc'):
            filename = filename[:-1]
        
        ast = compiler.parseFile(filename)
        walker = compiler.visitor.ASTVisitor()
        visitor = ExtractFunctionsVisitor([f.__name__ for f in functions])
        walker.preorder(ast, visitor)
        
        for func in functions:
            try:
                func.__ast__ = visitor.collectedFunctions[func.__name__]
            except KeyError:
                raise 'unable to find function "'+func.__name__+'" in "'+filename+'"'
        
    for c in compositions:
        c.optimize()

compositions = []
composed_functions = set()

class ComposedTransducer(object):
    def __init__(self, downstream, upstream):
        self.downstream = downstream
        self.upstream = upstream
        self.optimized = None
    
    def __call__(self):
        if self.optimized:
            exec self.optimized
        else:
            return self.downstream(self.upstream())
    
    def optimize(self):
        if self.optimized:
            return
        
        # this is where cycles would manifest themselves
        # TODO: break cycles somehow
        if type(self.upstream) == ComposedTransducer:
            print 'optimizing upstream'
            self.upstream.optimize()
        
        if type(self.downstream) == ComposedTransducer:
            print 'optimizing upstream'
            self.downstream.optimize()
        
        if hasattr(self.upstream,'__ast__'):
            print self.upstream.__ast__
        if hasattr(self.downstream,'__ast__'):
            print self.downstream.__ast__
        
        
        # build some ast
        #compiler.misc.set_filename(filename+'_composed_transducers', ast)
        #self.optimized = compiler.pycodegen.ModuleCodeGenerator(ast).getCode()
        pass


# internal

class ExtractFunctionsVisitor(object):
    def __init__(self, names):
        self.names = set(names)
        self.collectedFunctions = {}
    
    def visitFunction(self, node, *args):
        if node.name in self.names:
            name = node.name
            assert name not in self.collectedFunctions
            self.collectedFunctions[name] = node
        
        for child in node.getChildNodes():
            self.visit(child,*args)

class ComposeTransducersVisitor(object):
    def visitFunction(self, node, *args):
        print node.name
        node.name = node.name+'_exit_context'
        for child in node.getChildNodes():
            self.visit(child,*args)
        
        if not node.code.hasReturns:
            node.code.nodes.append(returnLocals)
